from PyQt5.QtWidgets import QOpenGLWidget
from PyQt5.QtGui import QSurfaceFormat, QVector3D, QMatrix4x4, QColor
from PyQt5.QtCore import QTimer
import OpenGL.GL as gl
import math

class GLWidget(QOpenGLWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.pose = {
            'x': 0.0,
            'y': 0.0,
            'z': 0.0,
            'roll': 0.0,
            'pitch': 0.0,
            'yaw': 0.0
        }
        self.camera_distance = -10.0  # 相机距离
        self.rotation_matrix = QMatrix4x4()
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update)
        self.timer.start(16)  # ~60 FPS
        self.setFormat(QSurfaceFormat.defaultFormat())

    def initializeGL(self):
        """初始化OpenGL环境"""
        gl.glClearColor(0.15, 0.15, 0.15, 1.0)
        gl.glEnable(gl.GL_DEPTH_TEST)
        gl.glEnable(gl.GL_LIGHTING)
        gl.glEnable(gl.GL_LIGHT0)
        gl.glEnable(gl.GL_COLOR_MATERIAL)
        gl.glColorMaterial(gl.GL_FRONT_AND_BACK, gl.GL_AMBIENT_AND_DIFFUSE)
        
        # 设置光源
        light_pos = [1.0, 1.0, 1.0, 0.0]
        gl.glLightfv(gl.GL_LIGHT0, gl.GL_POSITION, light_pos)

    def paintGL(self):
        """绘制场景"""
        gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
        gl.glLoadIdentity()
        
        # 设置相机视角
        gl.glTranslatef(0.0, 0.0, self.camera_distance)
        gl.glMultMatrixd(self.rotation_matrix.data())
        
        # 绘制坐标系
        self.draw_coordinate_system()
        
        # 绘制无人机模型
        self.draw_drone()

    def resizeGL(self, width, height):
        """调整视口"""
        aspect = width / height if height != 0 else 1
        gl.glViewport(0, 0, width, height)
        gl.glMatrixMode(gl.GL_PROJECTION)
        gl.glLoadIdentity()
        gl.glFrustum(-aspect, aspect, -1.0, 1.0, 5.0, 100.0)
        gl.glMatrixMode(gl.GL_MODELVIEW)

    def update_pose(self, pose):
        """更新无人机姿态"""
        self.pose.update(pose)
        self.update_rotation_matrix()
        self.update()

    def update_rotation_matrix(self):
        """根据欧拉角更新旋转矩阵"""
        self.rotation_matrix.setToIdentity()
        self.rotation_matrix.rotate(self.pose['roll'], 1, 0, 0)
        self.rotation_matrix.rotate(self.pose['pitch'], 0, 1, 0)
        self.rotation_matrix.rotate(self.pose['yaw'], 0, 0, 1)
        
    def draw_coordinate_system(self):
        """绘制坐标系轴线"""
        length = 2.0
        gl.glDisable(gl.GL_LIGHTING)
        
        # X轴（红色）
        gl.glColor3f(1.0, 0.0, 0.0)
        gl.glBegin(gl.GL_LINES)
        gl.glVertex3f(0.0, 0.0, 0.0)
        gl.glVertex3f(length, 0.0, 0.0)
        gl.glEnd()
        
        # Y轴（绿色）
        gl.glColor3f(0.0, 1.0, 0.0)
        gl.glBegin(gl.GL_LINES)
        gl.glVertex3f(0.0, 0.0, 0.0)
        gl.glVertex3f(0.0, length, 0.0)
        gl.glEnd()
        
        # Z轴（蓝色）
        gl.glColor3f(0.0, 0.0, 1.0)
        gl.glBegin(gl.GL_LINES)
        gl.glVertex3f(0.0, 0.0, 0.0)
        gl.glVertex3f(0.0, 0.0, length)
        gl.glEnd()
        
        gl.glEnable(gl.GL_LIGHTING)

    def draw_drone(self):
        """绘制无人机模型"""
        # 主体
        gl.glPushMatrix()
        gl.glColor3f(0.2, 0.2, 0.2)
        gl.glScalef(1.0, 0.2, 0.2)
        self.draw_cube()
        gl.glPopMatrix()
        
        # 旋翼臂
        for angle in [0, 90, 180, 270]:
            gl.glPushMatrix()
            gl.glRotatef(angle, 0, 0, 1)
            gl.glTranslatef(0.8, 0.0, 0.0)
            gl.glColor3f(0.5, 0.5, 0.5)
            gl.glScalef(0.4, 0.05, 0.05)
            self.draw_cube()
            gl.glPopMatrix()
            
            # 电机
            gl.glPushMatrix()
            gl.glRotatef(angle, 0, 0, 1)
            gl.glTranslatef(1.0, 0.0, 0.0)
            gl.glColor3f(0.8, 0.8, 0.8)
            gl.glScalef(0.2, 0.2, 0.2)
            self.draw_cube()
            gl.glPopMatrix()

    def draw_cube(self):
        """绘制立方体"""
        vertices = [
            [1, 1, -1], [-1, 1, -1], [-1, 1, 1], [1, 1, 1],
            [1, -1, -1], [-1, -1, -1], [-1, -1, 1], [1, -1, 1]
        ]
        edges = [
            (0, 1), (1, 2), (2, 3), (3, 0),
            (4, 5), (5, 6), (6, 7), (7, 4),
            (0, 4), (1, 5), (2, 6), (3, 7)
        ]
        
        gl.glBegin(gl.GL_LINES)
        for edge in edges:
            for vertex in edge:
                gl.glVertex3fv(vertices[vertex])
        gl.glEnd()